"""
自定义工具的封装。
但：参数是定好的，只能改写描述。
"""
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from pydantic import BaseModel, Field

## 参数是定好的，只能改写描述。
class WikipediaInputs(BaseModel):
    """维基百科工具的输入"""
    query: str = Field(
        description="query to look up in wikipedia,should be 3 or less words"
    )


## top_k_results=1: 只要一个结果，doc_content_chars_max=100 字数最大100
api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=100)

tool = WikipediaQueryRun(
    name="wikipedia-tool",
    description="look up things in wikipedia",
    args_schema=WikipediaInputs,
    api_wrapper=api_wrapper,
)

print(tool.run("langchain"))
